package mr

import (
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/rpc"
	"os"
	"time"
)

type TaskStatus int

const (
	Pending TaskStatus = iota
	Running
	Finished
)

func (s TaskStatus) String() string {
	switch s {
	case Pending:
		return "Pending"
	case Running:
		return "Running"
	case Finished:
		return "Finished"
	default:
		return fmt.Sprintf("%d", s)
	}
}

type MapTask struct {
	status TaskStatus
	path   string
	data   string
}

type ReduceTask struct {
	status TaskStatus
	paths  []string
}

type LocalOp struct {
	op   interface{}
	done chan bool
}

type Master struct {
	numMap int
	maps   []MapTask
	mIndex int
	mChan  chan LocalOp

	numReduce int
	reduces   []ReduceTask
	rIndex    int
	rChan     chan LocalOp
}

// some goroutines are dedicated to serve querys and updates to the
// master struct.

/**
 * stats transitions:
 * Pending (get)-> Running (commit)-> Finished
 * Running/Pending (reset)-> Pending
 */

type GetTask struct {
	any   bool
	index int
	task  interface{}
}

// obtain a pending map task by scanning `m.maps`.
// the status of the task is set to running before returning to the
// caller.
func (m *Master) getMapTask(r *GetTask) {
	r.any = false

	for j := 0; j < len(m.maps); j++ {
		i := m.mIndex
		t := m.maps[i]
		m.mIndex++
		if m.mIndex >= len(m.maps) {
			m.mIndex = 0
		}

		if t.status == Pending {
			m.maps[i].status = Running
			*r = GetTask{any: true, index: i, task: &t}
			break
		}
	}
}

func (m *Master) getReduceTask(r *GetTask) {
	r.any = false

	for j := 0; j < len(m.reduces); j++ {
		i := m.rIndex
		t := m.reduces[i]
		m.rIndex++
		if m.rIndex >= len(m.reduces) {
			m.rIndex = 0
		}

		if t.status == Pending {
			m.reduces[i].status = Running
			*r = GetTask{any: true, index: i, task: &t}
			break
		}
	}
}

type CommitTask struct {
	index int
	ok    bool
}

// mark a map task as finished.
func (m *Master) commitMapTask(r *CommitTask) {
	t := &m.maps[r.index]
	if t.status == Running {
		t.status = Finished
		r.ok = true
	} else {
		r.ok = false
	}
}

func (m *Master) commitReduceTask(r *CommitTask) {
	t := &m.reduces[r.index]
	if t.status == Running {
		t.status = Finished
		r.ok = true
	} else {
		r.ok = false
	}
}

type ResetTask struct {
	index int
	ok    bool
}

// reset a map task's status to pending if it's not finished.
func (m *Master) resetMapTask(r *ResetTask) {
	t := &m.maps[r.index]
	if t.status != Finished {
		log.Printf("map task #%d has been reset.", r.index)
		t.status = Pending
		r.ok = true
	} else {
		r.ok = false
	}
}

func (m *Master) resetReduceTask(r *ResetTask) {
	t := &m.reduces[r.index]
	if t.status != Finished {
		log.Printf("reduce task #%d has been reset.", r.index)
		t.status = Pending
		r.ok = true
	} else {
		r.ok = false
	}
}

// operation server functions

func (m *Master) mapChServer() {
	for req := range m.mChan {
		switch op := req.op.(type) {
		case *GetTask:
			m.getMapTask(op)
		case *CommitTask:
			m.commitMapTask(op)
		case *ResetTask:
			m.resetMapTask(op)

		default:
			log.Fatalf("unknown task op type: %T", op)
		}

		if req.done != nil {
			req.done <- true
		}
	}
}

func (m *Master) reduceChServer() {
	for req := range m.rChan {
		switch op := req.op.(type) {
		case *GetTask:
			m.getReduceTask(op)
		case *CommitTask:
			m.commitReduceTask(op)
		case *ResetTask:
			m.resetReduceTask(op)

		default:
			log.Fatalf("unknown intermediate op type: %T", op)
		}

		if req.done != nil {
			req.done <- true
		}
	}
}

// Your code here -- RPC handlers for the worker to call.

func taskTimeout(ch chan<- LocalOp, i int) {
	time.Sleep(10 * time.Second)
	u := ResetTask{index: i}
	ch <- LocalOp{op: &u, done: nil}
}

func (m *Master) FetchTask(args *FetchTaskArgs, reply *FetchTaskReply) error {
	u := GetTask{}
	done := make(chan bool)
	op := LocalOp{op: &u, done: done}

	reply.Any = false

	// find a map task first
	m.mChan <- op
	<-done

	if u.any {
		reply.Any = true
		reply.Kind = TaskKindMap
		reply.Index = u.index
		go taskTimeout(m.mChan, u.index)
		return nil
	}

	// then try to find a reduce task
	m.rChan <- op
	<-done

	if u.any {
		reply.Any = true
		reply.Kind = TaskKindReduce
		reply.Index = u.index
		go taskTimeout(m.rChan, u.index)
		return nil
	}

	return nil
}

func (m *Master) FetchMapTask(args *FetchMapTaskArgs, reply *FetchMapTaskReply) error {
	i := args.Index
	if i < 0 || i >= len(m.maps) {
		return fmt.Errorf("index out of range")
	}

	t := m.maps[i]
	reply.Path = t.path
	reply.Data = t.data
	reply.NumReduce = m.numReduce

	return nil
}

func (m *Master) FetchReduceTask(args *FetchReduceTaskArgs, reply *FetchReduceTaskReply) error {
	i := args.Index
	if i < 0 || i >= len(m.reduces) {
		return fmt.Errorf("index out of range")
	}

	t := m.reduces[i]
	reply.Paths = t.paths
	reply.NumMap = m.numMap

	return nil
}

func (m *Master) ReportMapTask(args *ReportMapTaskArgs, reply *ReportMapTaskReply) error {
	done := make(chan bool)
	u := CommitTask{index: args.Index}
	m.mChan <- LocalOp{op: &u, done: done}
	<-done

	if u.ok {
		log.Printf("received intermediate files: %v", args.Paths)
		for j, path := range args.Paths {
			m.reduces[j].paths = append(m.reduces[j].paths, path)
		}
	}

	return nil
}

func (m *Master) ReportReduceTask(args *ReportReduceTaskArgs, reply *ReportReduceTaskReply) error {
	u := CommitTask{index: args.Index}
	m.rChan <- LocalOp{op: &u, done: nil}
	return nil
}

//
// start a thread that listens for RPCs from worker.go
//
func (m *Master) startRPCServer() {
	rpc.Register(m)
	rpc.HandleHTTP()
	//l, e := net.Listen("tcp", ":1234")
	sockname := masterSock()
	os.Remove(sockname)
	l, e := net.Listen("unix", sockname)
	if e != nil {
		log.Fatal("listen error:", e)
	}
	go http.Serve(l, nil)
}

//
// main/mrmaster.go calls Done() periodically to find out
// if the entire job has finished.
//
func (m *Master) Done() bool {
	for _, t := range m.reduces {
		if t.status != Finished {
			return false
		}
	}
	return true
}

// load file contents specified in paths.
func (m *Master) loadTasks(paths []string) {
	ch := make(chan MapTask)

	for _, path := range paths {
		go func(path string) {
			log.Printf("loading %q...", path)
			fp, err := os.Open(path)
			if err != nil {
				log.Fatalf("cannot open file %q.", path)
			}
			defer fp.Close()

			data, err := io.ReadAll(fp)
			if err != nil {
				log.Fatalf("cannot read file %q.", path)
			}

			ch <- MapTask{status: Pending, path: path, data: string(data)}
		}(path)
	}

	m.maps = make([]MapTask, m.numMap)
	for i := range paths {
		m.maps[i] = <-ch
	}
}

//
// create a Master.
// main/mrmaster.go calls this function.
// nReduce is the number of reduce tasks to use.
//
func MakeMaster(files []string, nReduce int) *Master {
	setupLog("master")

	m := Master{}
	m.numMap = len(files)
	m.numReduce = nReduce

	m.loadTasks(files)

	m.reduces = make([]ReduceTask, nReduce)
	for i := 0; i < nReduce; i++ {
		m.reduces[i].paths = make([]string, 0)
	}

	m.mChan = make(chan LocalOp)
	go m.mapChServer()

	m.rChan = make(chan LocalOp)
	go m.reduceChServer()

	m.startRPCServer()

	return &m
}
